#include <Memory/Heap.h>
#include <Memory/Paging.h>
#include <IntExcept/Debug.h>
#include <Utility/intrinsic.h>

/*
@brief 释放堆内存单元
@param `lpHeap` 堆地址
@param `lpAddress` 待释放的地址
@return 状态码
@note 优化方案：遍历时按照地址对堆进行排序
*/
DWORD HeapFree(LPHeap lpHeap, LPVOID lpAddress)
{
    /* 堆中的 Arena 描述符存放了起始地址，每个 Arena 占用 1 页，因此非常容易确定释放的地址范围 */
    LPVOID lpArena = (LPVOID)((DWORD_PTR)lpAddress & 0xFFFFF000);
    int nArnDscp = 0;
    /* 定位到 lpAddress 相应的 Arena 描述符 */
    while (TRUE)
    {
        if (lpHeap[nArnDscp].wGranularity != ARENA_NONE && lpHeap[nArnDscp].lpArena == lpArena)
        {
            break;
        }
        nArnDscp++;
    }
    /* 所有已申请的内存块其对应位图中的位必定清零，可据此判断释放操作是否合法 */
    if (lpHeap[nArnDscp].wGranularity < 256) /* 小型 Arena */
    {
        /* 位图记录与描述符中记录必须一致 */
        WORD wAvailable = 0;
        for (int j = 0; j < ARENA_BITMAP_SIZE_IN_DWORD(lpHeap[nArnDscp].wGranularity); j++)
        {
            wAvailable += __popcnt(((LPDWORD)(lpHeap[nArnDscp].lpArena))[j]);
        }
        KERNEL_ASSERT(wAvailable == lpHeap[nArnDscp].wAvailableUnits, "Number of bits available in bitmap inconsistent with wAvailableUnits in arena descriptor.\r\n");
        if (lpHeap[nArnDscp].wGranularity == 16)
        {
            KERNEL_ASSERT((*(LPDWORD)(lpHeap[nArnDscp].lpArena) & 0x3) == 0, "First two bits of 16 B arena bitmap non-zero.\r\n"); /* 16 B Arena 位图首 2 位必为 0 */
            DWORD dwOffset = ((DWORD_PTR)lpAddress - (DWORD_PTR)lpArena) >> 4; /* 位图中的偏移量 */
            KERNEL_ASSERT(dwOffset > 1 && dwOffset != MEMORY_BITMAP_INVALID_OFFSET, "Physical memory pool runs off.\r\n"); /* 位图偏移量必须大于 1，越过位图本身所占空间 */
            KERNEL_ASSERT(MEMORY_BITMAP_INVALID_OFFSET != ReleaseOccupiedBit(dwOffset, lpHeap[nArnDscp].lpArena), "Attempt to release illegal bit in arena bitmap.\r\n");
        }
        else 
        {
            KERNEL_ASSERT((*(LPDWORD)(lpHeap[nArnDscp].lpArena) & 0x1) == 0, "First bit of arena bitmap non-zero.\r\n"); /* 其余种类的小型 Arena 位图第一位必为 0 */
            DWORD dwOffset = ((DWORD_PTR)lpAddress - (DWORD_PTR)lpArena) >> BIT_SCAN_FORWARD_DWORD(lpHeap[nArnDscp].wGranularity); /* 位图中的偏移量 */
            KERNEL_ASSERT(dwOffset > 0 && dwOffset != MEMORY_BITMAP_INVALID_OFFSET, "Physical memory pool runs off.\r\n");
            KERNEL_ASSERT(MEMORY_BITMAP_INVALID_OFFSET != ReleaseOccupiedBit(dwOffset, lpHeap[nArnDscp].lpArena), "Attempt to release illegal bit in arena bitmap.\r\n");
        }
        lpHeap[nArnDscp].wAvailableUnits++;
    }
    else if (lpHeap[nArnDscp].wGranularity < 4096) /* 中型 Arena */
    {
        DWORD dwOffset = ((DWORD_PTR)lpAddress - (DWORD_PTR)lpArena) >> BIT_SCAN_FORWARD_DWORD(lpHeap[nArnDscp].wGranularity); /* 位图中的偏移量 */
        DWORD dwBitmap = lpHeap[nArnDscp].wAvailableUnits;
        KERNEL_ASSERT(MEMORY_BITMAP_INVALID_OFFSET != ReleaseOccupiedBit(dwOffset, &dwBitmap), "Attempt to release illegal bit in arena bitmap\r\n");
        lpHeap[nArnDscp].wAvailableUnits = (WORD)dwBitmap;
    }
    else /* 大型 Arena */
    {
        KERNEL_ASSERT(lpAddress == lpHeap[nArnDscp].lpArena, "Attempt to release unknown memory block.\r\n");
        KERNEL_ASSERT(PAGE_FREE_SUCCESS == FreePage(lpAddress, lpHeap[nArnDscp].wAvailableUnits), "Error while trying to free virtual page.\r\n"); /* 释放所有页框 */
        lpHeap[nArnDscp].wGranularity = ARENA_NONE;
    }
    return 0;
}